# -*- coding: utf-8 -*-
# @Time    : 2023/6/1 5:13 下午
# @Author  : Wu WanJie
import torch
import torch.nn as nn
import torch.nn.functional as F


class FocalLoss(nn.Module):
    def __init__(self, alpha=120, gamma=2, reduction='mean'):
        super(FocalLoss, self).__init__()
        self.alpha = alpha
        self.gamma = gamma
        self.reduction = reduction

    def forward(self, inputs, targets):
        # 输入的shape：(batch, num_classes)
        # 目标的shape：(batch,)

        # 计算softmax和交叉熵损失
        ce_loss = F.cross_entropy(inputs, targets, reduction='none')

        # 计算 p_t
        p_t = torch.exp(-ce_loss)

        # 计算 focal loss
        focal_loss = self.alpha * (1 - p_t) ** self.gamma * ce_loss

        # 对损失进行归约
        if self.reduction == 'mean':
            return torch.mean(focal_loss)
        elif self.reduction == 'sum':
            return torch.sum(focal_loss)
        else:
            return focal_loss
